{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "LatentConstraints_magenta.ipynb",
      "version": "0.3.2",
      "views": {},
      "default_view": {},
      "provenance": [
        {
          "file_id": "1GPrBZWol4z_5R-LHmOjX3Y8YBdQ-92Rr",
          "timestamp": 1510266173168
        },
        {
          "file_id": "1-RpoOTJxM0AzFskOwW7qXB58sBe4HQJC",
          "timestamp": 1510185306485
        },
        {
          "file_id": "0BxV6RES_0xkhSE5nUTRQbldqTmc",
          "timestamp": 1506463290551
        }
      ],
      "collapsed_sections": [],
      "toc_visible": true
    }
  },
  "cells": [
    {
      "metadata": {
        "id": "oKI30KBUEN6l",
        "colab_type": "text"
      },
      "source": [
        "Copyright 2017 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "8w2CSFQvkALu",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "A version of this notebook with accompanying cloud instance can be found at [https://colab.research.google.com/notebook#fileId=/v2/external/notebooks/latent_constraints/latentconstraints.ipynb](https://colab.research.google.com/notebook#fileId=/v2/external/notebooks/latent_constraints/latentconstraints.ipynb&scrollTo=oKI30KBUEN6l)\n",
        "\n",
        "# Latent Constraints: Conditional Generation from Unconditional Generative Models\n",
        "### ___Jesse Engel, Matthew Hoffman, Adam Roberts___ [arXiv link](http://arxiv.org/abs/1711.05772)\n",
        "\n",
        "___\n",
        "__Abstract:__ Deep generative neural networks have proven effective at both conditional and unconditional modeling of complex data distributions.\n",
        "Conditional generation enables interactive control, but creating new controls often requires expensive retraining. \n",
        "In this paper, we develop a method to condition generation without retraining the model. \n",
        "By post-hoc learning ___latent constraints___, value functions that identify regions in latent space that generate outputs with desired attributes, we can conditionally sample from these regions with gradient-based optimization or amortized actor functions.\n",
        "Combining attribute constraints with a universal ``realism`` constraint, which enforces similarity to the data distribution, we generate realistic conditional images from an unconditional variational autoencoder.\n",
        "Further, using gradient-based optimization, we demonstrate identity-preserving transformations that make\n",
        "the minimal adjustment in latent space to modify the attributes of an image.\n",
        "Finally, with discrete sequences of musical notes, we demonstrate zero-shot conditional generation, learning latent constraints in the absence of labeled data or a differentiable reward function. \n",
        "\n",
        "___\n",
        "\n",
        "![](http://download.magenta.tensorflow.org/models/latent_constraints/ipynb_figs/diagram.png =400x400 \"Latent Space Diagram\")\n",
        "![](http://download.magenta.tensorflow.org/models/latent_constraints/ipynb_figs/cgen.png =540x500 \"Conditional Generation with and without a distance penalty\")\n",
        "\n",
        "___\n",
        "\n",
        "This notebook contains code for running experiments related to the paper. First, we load pretrained checkpoints:\n",
        "\n",
        "  * VAE models trained on CelebA with pixelwise gaussian data liklihoods of $\\mathcal{N}(\\mu(z), \\sigma_x=0.1)$ and $\\mathcal{N}(\\mu(z), \\sigma_x=1)$. \n",
        "  * We also provide embeddings of the training and eval set from the VAE models.\n",
        "  * A generator ($G$) and discriminator ($D$) from a conditional-GAN, trained to shift samples from the prior to new points in latent space that satisfy the realism constraint ($r$) and attribute constraints ($r_{attr}$). \n",
        "  * We have versions trained with no distance penalty, and also with a penalty of 1e-1.\n",
        "  * A seperately trained attribute classifier in both z-space ($D_{attr}$) and pixel space ($Classifier$).\n",
        "\n",
        "We then proceed to:\n",
        " * Demonstrate that VAE reconstructions sharpen as $\\sigma_x$ lowers, at the expense of sample quality, which is compensated with latent constraints.\n",
        " * Plot conditional generation using CGANs ($D$, $G$) both with and without distance penalty.\n",
        " * Perform identity preserving transformations doing SGD in z-space wrt $D_{attr}$\n",
        " * Evaluate the accuracy of generating images with conditional attributes.\n",
        " \n",
        "Training loops are also provided for demonstration purposes at the end of the notebook.\n",
        "\n",
        "___\n",
        "\n",
        "This colab notebook is self-contained and should run natively on google cloud. The code and checkpoints can be downloaded separately and run locally, which is recommended if you want to train your own model. Pretrained model checkpoints are available at [download.magenta.tensorflow.org/models/latent_constraints/latent_constraints.tar](http://download.magenta.tensorflow.org/models/latent_constraints/latent_constraints.tar). \n",
        "\n",
        "Download them and extract files to `/tmp/` where this notebook assumes that the files exist.\n",
        "\n",
        "__Tips:__ Don't forget you can navigate with the _Table of Contents_ in the left hand sidebar, and collapse all sections with (`Ctrl + Shift + ]`)\n",
        "\n",
        "\n"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "7PMgN4pHd9XD",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# This notebook requires DeepMind's sonnet library, which itself\n",
        "# requires the nightly build of TensorFlow. The command below \n",
        "# installs both.\n",
        "!pip install -q -U dm-sonnet tf-nightly\n",
        "\n",
        "import os\n",
        "import PIL\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import sklearn.metrics\n",
        "import sonnet as snt\n",
        "import tensorflow as tf\n",
        "\n",
        "ds = tf.contrib.distributions\n",
        "\n",
        "%matplotlib inline"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "PNHtTJ74wsHM",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Load the Data"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "8h8M3idWxLb5",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "basepath = '/tmp/'\n",
        "\n",
        "# Load CelebA embeddings\n",
        "# VAE with x_sigma = 0.1\n",
        "train_mu = np.load(basepath + 'train_mu.npy')\n",
        "train_sigma = np.load(basepath + 'train_sigma.npy')\n",
        "eval_mu = np.load(basepath + 'eval_mu.npy')\n",
        "eval_sigma = np.load(basepath + 'eval_sigma.npy')\n",
        "\n",
        "# VAE with x_sigma = 1.0\n",
        "eval_mu_xsigma1 = np.load(basepath + 'eval_mu_xsigma1.npy')\n",
        "eval_sigma_xsigma1 = np.load(basepath + 'eval_sigma_xsigma1.npy')\n",
        "\n",
        "np.random.seed(10003)\n",
        "n_train = train_mu.shape[0]\n",
        "n_eval = eval_mu.shape[0]\n",
        "\n",
        "# Load Attributes\n",
        "# Only use 10 salient attributes\n",
        "attr_train = np.load(basepath + 'attr_train.npy')\n",
        "attr_eval = np.load(basepath + 'attr_eval.npy')\n",
        "attr_test = np.load(basepath + 'attr_test.npy')\n",
        "\n",
        "attr_mask = [4, 8, 9, 11, 15, 20, 24, 31, 35, 39]\n",
        "attribute_names = [\n",
        "    'Bald',\n",
        "    'Black_Hair',\n",
        "    'Blond_Hair',\n",
        "    'Brown_Hair',\n",
        "    'Eyeglasses',\n",
        "    'Male',\n",
        "    'No_Beard',\n",
        "    'Smiling',\n",
        "    'Wearing_Hat',\n",
        "    'Young',\n",
        "]\n",
        "\n",
        "attr_train = attr_train[:, attr_mask]\n",
        "attr_eval = attr_eval[:, attr_mask]\n",
        "attr_test = attr_test[:, attr_mask]"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "948GFxky-F7I",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Define the Graph\n",
        "\n",
        "All the functions with variables are wrapped in `sonnet` modules. \n",
        "\n",
        "As is the ``model`` that ties them all together. \n",
        "\n",
        "Tensors (endpoints) are accessible as attributes of the ``model``. "
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "M11bgT4orksE",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "class Encoder(snt.AbstractModule):\n",
        "  '''VAE Convolutional Encoder.'''\n",
        "  def __init__(self,\n",
        "               n_latent,\n",
        "               layers=((256, 5, 2),\n",
        "                       (512, 5, 2),\n",
        "                       (1024, 3, 2),\n",
        "                       (2048, 3, 2)),\n",
        "               name='encoder'):\n",
        "    super(Encoder, self).__init__(name=name)\n",
        "    self.n_latent = n_latent\n",
        "    self.layers = layers\n",
        "\n",
        "  def _build(self, x):\n",
        "    h = x\n",
        "    for unused_i, l in enumerate(self.layers):\n",
        "      h = tf.nn.relu(snt.Conv2D(l[0], l[1], l[2])(h))\n",
        "\n",
        "    h_shape = h.get_shape().as_list()\n",
        "    h = tf.reshape(h, [-1, h_shape[1] * h_shape[2] * h_shape[3]])\n",
        "    pre_z = snt.Linear(2 * self.n_latent)(h)\n",
        "    mu = pre_z[:, :self.n_latent]\n",
        "    sigma = tf.nn.softplus(pre_z[:, self.n_latent:])\n",
        "    return mu, sigma\n",
        "\n",
        "\n",
        "class Decoder(snt.AbstractModule):\n",
        "  '''VAE Convolutional Decoder.'''\n",
        "  def __init__(self,\n",
        "               layers=((2048, 4, 4),\n",
        "                       (1024, 3, 2),\n",
        "                       (512, 3, 2),\n",
        "                       (256, 5, 2),\n",
        "                       (3, 5, 2)),\n",
        "               name='decoder'):\n",
        "    super(Decoder, self).__init__(name=name)\n",
        "    self.layers = layers\n",
        "\n",
        "  def _build(self, x):\n",
        "    for i, l in enumerate(self.layers):\n",
        "      if i == 0:\n",
        "        h = snt.Linear(l[1] * l[2] * l[0])(x)\n",
        "        h = tf.reshape(h, [-1, l[1], l[2], l[0]])\n",
        "      elif i == len(self.layers) - 1:\n",
        "        h = snt.Conv2DTranspose(l[0], None, l[1], l[2])(h)\n",
        "      else:\n",
        "        h = tf.nn.relu(snt.Conv2DTranspose(l[0], None, l[1], l[2])(h))\n",
        "    logits = h\n",
        "    return logits\n",
        "\n",
        "\n",
        "class G(snt.AbstractModule):\n",
        "  '''CGAN Generator. Maps from z-space to z-space.'''\n",
        "  def __init__(self,\n",
        "               n_latent,\n",
        "               layers=(2048,)*4,\n",
        "               name='generator'):\n",
        "    super(G, self).__init__(name=name)\n",
        "    self.layers = layers\n",
        "    self.n_latent = n_latent\n",
        "\n",
        "  def _build(self, z_and_labels):\n",
        "    z, labels = z_and_labels\n",
        "    labels = tf.cast(labels, tf.float32)\n",
        "    size = self.layers[0]\n",
        "    x = tf.concat([z, snt.Linear(size)(labels)], axis=-1)\n",
        "    for l in self.layers:\n",
        "      x = tf.nn.relu(snt.Linear(l)(x))\n",
        "    x = snt.Linear(2 * self.n_latent)(x)\n",
        "    dz = x[:, :self.n_latent]\n",
        "    gates = tf.nn.sigmoid(x[:, self.n_latent:])\n",
        "    z_prime = (1-gates) * z + gates * dz\n",
        "    return z_prime\n",
        "\n",
        "\n",
        "class D(snt.AbstractModule):\n",
        "  '''CGAN Discriminator.'''\n",
        "  def __init__(self,\n",
        "               output_size=1,\n",
        "               layers=(2048,)*4,\n",
        "               name='D'):\n",
        "    super(D, self).__init__(name=name)\n",
        "    self.layers = layers\n",
        "    self.output_size = output_size\n",
        "\n",
        "  def _build(self, z_and_labels):\n",
        "    z, labels = z_and_labels\n",
        "    labels = tf.cast(labels, tf.float32)\n",
        "    size = self.layers[0]\n",
        "    x = tf.concat([z, snt.Linear(size)(labels)], axis=-1)\n",
        "    for l in self.layers:\n",
        "      x = tf.nn.relu(snt.Linear(l)(x))\n",
        "    logits = snt.Linear(self.output_size)(x)\n",
        "    return logits\n",
        "\n",
        "\n",
        "class DAttr(snt.AbstractModule):\n",
        "  '''Attribute Classifier from z-space.'''\n",
        "  def __init__(self,\n",
        "               output_size=1,\n",
        "               layers=(2048,)*4,\n",
        "               name='D'):\n",
        "    super(DAttr, self).__init__(name=name)\n",
        "    self.layers = layers\n",
        "    self.output_size = output_size\n",
        "\n",
        "  def _build(self, x):\n",
        "    for l in self.layers:\n",
        "      x = tf.nn.relu(snt.Linear(l)(x))\n",
        "    logits = snt.Linear(self.output_size)(x)\n",
        "    return logits\n",
        "  \n",
        "  \n",
        "class Classifier(snt.AbstractModule):\n",
        "  '''Convolutional Attribute Classifier from Pixels.'''\n",
        "  def __init__(self,\n",
        "               output_size,\n",
        "               layers=((256, 5, 2),\n",
        "                       (256, 3, 1),\n",
        "                       (512, 5, 2),\n",
        "                       (512, 3, 1),\n",
        "                       (1024, 3, 2),\n",
        "                       (2048, 3, 2)),\n",
        "               name='encoder'):\n",
        "    super(Classifier, self).__init__(name=name)\n",
        "    self.output_size = output_size\n",
        "    self.layers = layers\n",
        "\n",
        "  def _build(self, x):\n",
        "    h = x\n",
        "    for unused_i, l in enumerate(self.layers):\n",
        "      h = tf.nn.relu(snt.Conv2D(l[0], l[1], l[2])(h))\n",
        "\n",
        "    h_shape = h.get_shape().as_list()\n",
        "    h = tf.reshape(h, [-1, h_shape[1] * h_shape[2] * h_shape[3]])\n",
        "    logits = snt.Linear(self.output_size)(h)\n",
        "    return logits"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yGFCiG1tr9Bi",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "class Model(snt.AbstractModule):\n",
        "  '''All the components glued together.'''\n",
        "  def __init__(self, config, name=''):\n",
        "    super(Model, self).__init__(name=name)\n",
        "    self.config = config\n",
        "\n",
        "  def _build(self, unused_input=None):\n",
        "    config = self.config\n",
        "\n",
        "    # Constants\n",
        "    batch_size = config['batch_size']\n",
        "    n_latent = config['n_latent']\n",
        "    img_width = config['img_width']\n",
        "    half_batch = int(batch_size / 2)\n",
        "    n_labels = 10\n",
        "\n",
        "    #---------------------------------------------------------------------------\n",
        "    ### Placeholders\n",
        "    #---------------------------------------------------------------------------\n",
        "    x = tf.placeholder(tf.float32, \n",
        "                       shape=(None, img_width, img_width, 3), name='x')\n",
        "    # Attributes\n",
        "    labels = tf.placeholder(tf.int32, shape=(None, n_labels), name='labels')\n",
        "    # Real / fake label reward\n",
        "    r = tf.placeholder(tf.float32, shape=(None, 1), name='D_label')\n",
        "    # Transform through optimization\n",
        "    z0 = tf.placeholder(tf.float32, shape=(None, n_latent), name='z0')\n",
        "    z_prime = tf.get_variable('z_prime', \n",
        "                              shape=(half_batch, n_latent), dtype=tf.float32)\n",
        "\n",
        "    #---------------------------------------------------------------------------\n",
        "    ### Modules with parameters\n",
        "    #---------------------------------------------------------------------------\n",
        "    encoder = Encoder(n_latent=n_latent, name='encoder')\n",
        "    decoder = Decoder(name='decoder')\n",
        "    g = G(n_latent=n_latent, name='generator')\n",
        "    d = D(output_size=1, name='d_z')\n",
        "    d_attr = DAttr(output_size=n_labels, name='d_attr')\n",
        "    classifier = Classifier(output_size=n_labels, name='classifier')\n",
        "\n",
        "\n",
        "    #---------------------------------------------------------------------------\n",
        "    ### VAE\n",
        "    #---------------------------------------------------------------------------\n",
        "    # Encode\n",
        "    mu, sigma = encoder(x)\n",
        "    q_z = ds.Normal(loc=mu, scale=sigma)\n",
        "\n",
        "    # Optimize / Amortize or feedthrough\n",
        "    q_z_sample = q_z.sample()\n",
        "\n",
        "    transform = tf.constant(False)\n",
        "    z = tf.cond(transform, lambda: z_prime, lambda: q_z_sample)\n",
        "\n",
        "    amortize = tf.constant(False)\n",
        "    z = tf.cond(amortize, lambda: g((z, labels)), lambda: z)\n",
        "\n",
        "    # Decode\n",
        "    logits = decoder(z)\n",
        "    x_sigma = tf.constant(config['x_sigma'])\n",
        "    p_x = ds.Normal(loc=tf.nn.sigmoid(logits), scale=x_sigma)\n",
        "    x_mean = p_x.mean()\n",
        "\n",
        "    # Reconstruction Loss\n",
        "    recons = tf.reduce_sum(p_x.log_prob(x), axis=[1, 2, 3])\n",
        "\n",
        "    mean_recons = tf.reduce_mean(recons)\n",
        "\n",
        "    # Prior\n",
        "    p_z = ds.Normal(loc=0., scale=1.)\n",
        "    prior_sample = p_z.sample(sample_shape=[batch_size, n_latent])\n",
        "\n",
        "    # KL Loss\n",
        "    KL_qp = ds.kl_divergence(q_z, p_z)\n",
        "    KL = tf.reduce_sum(KL_qp, axis=-1)\n",
        "    mean_KL = tf.reduce_mean(KL)\n",
        "\n",
        "    beta = tf.constant(config['beta'])\n",
        "\n",
        "    # VAE Loss\n",
        "    vae_loss = -mean_recons + mean_KL * beta\n",
        "\n",
        "    #---------------------------------------------------------------------------\n",
        "    ### Discriminator Constraint in Img and Z space and Digit space (implicit)\n",
        "    #---------------------------------------------------------------------------\n",
        "    d_logits = d([z, labels])\n",
        "\n",
        "    r_pred = tf.nn.sigmoid(d_logits)  # r = [0 prior, 1 data]\n",
        "    d_loss = tf.losses.sigmoid_cross_entropy(r, d_logits)\n",
        "\n",
        "    # Mean over examples\n",
        "    d_loss = tf.reduce_mean(d_loss)\n",
        "\n",
        "    # Gradient Penalty\n",
        "    real_data = z[:half_batch]\n",
        "    fake_data = z[half_batch:batch_size]\n",
        "    alpha = tf.random_uniform(shape=[half_batch, n_latent], minval=0., maxval=1.)\n",
        "    differences = fake_data - real_data\n",
        "    interpolates = real_data + (alpha * differences)\n",
        "    interp_pred = d([interpolates, labels[:half_batch]])\n",
        "    gradients = tf.gradients(interp_pred, [interpolates])[0]\n",
        "    slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]) + 1e-10)\n",
        "    gradient_penalty = tf.reduce_mean((slopes - 1.)**2)\n",
        "\n",
        "    # Add penalty\n",
        "    lambda_weight = tf.constant(config['lambda_weight'])\n",
        "    d_loss_training = d_loss + lambda_weight * gradient_penalty\n",
        "\n",
        "    \n",
        "    #---------------------------------------------------------------------------\n",
        "    ### Discriminator Attribute classification (implicit constraint)\n",
        "    #---------------------------------------------------------------------------\n",
        "    # Z-Space\n",
        "    attr_weights = tf.constant(np.ones([1, n_labels]).astype(np.float32))\n",
        "    logits_attr = d_attr(z)\n",
        "    pred_attr = tf.nn.sigmoid(logits_attr)\n",
        "    d_loss_attr =  tf.losses.sigmoid_cross_entropy(labels, \n",
        "                                                   logits=logits_attr, \n",
        "                                                   weights=attr_weights)\n",
        "\n",
        "    \n",
        "    #---------------------------------------------------------------------------\n",
        "    ### OPTIMIZTION TRANSFORMATION (SGD)\n",
        "    #---------------------------------------------------------------------------\n",
        "    # Realism Constraint\n",
        "    transform_r_weight = tf.constant(1.0)\n",
        "    loss_transform = transform_r_weight * tf.reduce_mean(d_loss)\n",
        "\n",
        "    # Attribute Constraint\n",
        "    transform_attr_weight = tf.constant(0.0)\n",
        "    loss_transform += transform_attr_weight * d_loss_attr\n",
        "    \n",
        "    # Distance Penalty\n",
        "    transform_penalty_weight = tf.constant(0.0)\n",
        "    z_sigma_mean = tf.constant(np.ones([1, n_latent]).astype(np.float32))\n",
        "    transform_penalty = tf.log(1 + (z_prime - z0)**2)\n",
        "    transform_penalty = transform_penalty * z_sigma_mean**-2\n",
        "    loss_transform += tf.reduce_mean(transform_penalty_weight * transform_penalty)\n",
        "\n",
        "\n",
        "    #---------------------------------------------------------------------------\n",
        "    ### AMORTIZED TRANSFORMATION (Generator)\n",
        "    #---------------------------------------------------------------------------\n",
        "    # Realism and Attribute Constraint\n",
        "    g_loss = -tf.log(tf.clip_by_value(r_pred, 1e-15, 1 - 1e-15))\n",
        "    g_loss = tf.reduce_mean(g_loss)\n",
        "\n",
        "    # Distance Penalty\n",
        "    g_penalty_weight = tf.constant(0.0)\n",
        "    g_penalty = tf.log(1 + (z - q_z_sample)**2)\n",
        "    g_penalty = g_penalty * z_sigma_mean**-2\n",
        "    g_penalty = tf.reduce_mean(g_penalty) \n",
        "    g_loss += g_penalty_weight * g_penalty\n",
        "\n",
        "    #---------------------------------------------------------------------------\n",
        "    ### Classify Attributes from pixels\n",
        "    #---------------------------------------------------------------------------\n",
        "    logits_classifier = classifier(x)\n",
        "    pred_classifier = tf.nn.sigmoid(logits_classifier)\n",
        "    classifier_loss =  tf.losses.sigmoid_cross_entropy(labels, \n",
        "                                                       logits=logits_classifier)\n",
        "\n",
        "    \n",
        "    #---------------------------------------------------------------------------\n",
        "    ### Training\n",
        "    #---------------------------------------------------------------------------\n",
        "    # Learning rates\n",
        "    d_lr = tf.constant(3e-4)\n",
        "    d_attr_lr = tf.constant(3e-4)\n",
        "    vae_lr = tf.constant(3e-4)\n",
        "    g_lr = tf.constant(3e-4)\n",
        "    classifier_lr = tf.constant(3e-4)\n",
        "    transform_lr = tf.constant(3e-4)\n",
        "\n",
        "    # Training Ops\n",
        "    vae_vars = list(encoder.get_variables())\n",
        "    vae_vars.extend(decoder.get_variables())\n",
        "    train_vae = tf.train.AdamOptimizer(vae_lr).minimize(vae_loss, var_list=vae_vars)\n",
        "\n",
        "    d_vars = d.get_variables()\n",
        "    train_d = tf.train.AdamOptimizer(d_lr, beta1=0, beta2=0.9).minimize(\n",
        "        d_loss_training, var_list=d_vars)\n",
        "\n",
        "    classifier_vars = classifier.get_variables()\n",
        "    train_classifier = tf.train.AdamOptimizer(classifier_lr).minimize(\n",
        "        classifier_loss, var_list=classifier_vars)\n",
        "\n",
        "    g_vars = g.get_variables()\n",
        "    train_g = tf.train.AdamOptimizer(g_lr, beta1=0, beta2=0.9).minimize(\n",
        "        g_loss, var_list=g_vars)\n",
        "\n",
        "    d_attr_vars = d_attr.get_variables()\n",
        "    train_d_attr = tf.train.AdamOptimizer(d_attr_lr).minimize(\n",
        "        d_loss_attr, var_list=d_attr_vars)\n",
        "\n",
        "    train_transform = tf.train.AdamOptimizer(transform_lr).minimize(\n",
        "            loss_transform, var_list=[z_prime])\n",
        "    \n",
        "    # Savers\n",
        "    vae_saver = tf.train.Saver(vae_vars, max_to_keep=100)\n",
        "    g_saver = tf.train.Saver(g_vars, max_to_keep=1000)\n",
        "    d_saver = tf.train.Saver(d_vars, max_to_keep=1000)\n",
        "    d_attr_saver = tf.train.Saver(d_attr_vars, max_to_keep=1000)\n",
        "    classifier_saver = tf.train.Saver(classifier_vars, max_to_keep=1000)\n",
        "\n",
        "    # Add all endpoints as object attributes\n",
        "    for k, v in locals().iteritems():\n",
        "      self.__dict__[k] = v"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "IRjPCq7OsbSq",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Load all models"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "uZ9Y8wqvLaCo",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "config = {\n",
        "    'n_latent': 1024,\n",
        "    'img_width': 64,\n",
        "    'crop_width': 64,\n",
        "    # Optimization parameters\n",
        "    'batch_size': 128,\n",
        "    'beta': 1.0,\n",
        "    'x_sigma': 0.1,\n",
        "    'lambda_weight': 10.0,\n",
        "    'penalty_weight': 0.0,\n",
        "}"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "haaFDnL5_qki",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "tf.reset_default_graph()\n",
        "sess = tf.Session()\n",
        "\n",
        "# Declare\n",
        "m = Model(config)\n",
        "# Build\n",
        "_ = m()\n",
        "# Initialize\n",
        "sess.run(tf.global_variables_initializer())"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sw82b_sx_XrS",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load VAE\n",
        "ckpt = os.path.join(basepath, 'vae_best_celeba_0_crop128_beta1.ckpt')\n",
        "m.vae_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "dSAqA5UWQ5pO",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load D\n",
        "ckpt = os.path.join(basepath, 'D_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.d_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "t9rpEbm9RFBv",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load G\n",
        "ckpt = os.path.join(basepath, 'G_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.g_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TZzIbNFFRHgP",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load D_attr\n",
        "ckpt = os.path.join(basepath, 'D_attr_best_d_attr_0.ckpt')\n",
        "m.d_attr_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "cicRZT2TRK4X",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load Classifier\n",
        "ckpt = os.path.join(basepath, 'classifier_best_classifier_0.ckpt')\n",
        "m.classifier_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_TCfRCP0Xcq5",
        "colab_type": "text"
      },
      "source": [
        "# GENERATE PLOTS"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "qRUeq0qIl6hG",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "def im(x):\n",
        "  plt.imshow(np.maximum(0, np.minimum(1, x)), interpolation='none')\n",
        "  plt.xticks([])\n",
        "  plt.yticks([])    \n",
        "  \n",
        "def batch_image(b, max_images=64, rows=None, cols=None):\n",
        "  \"\"\"Turn a batch of images into a single image mosaic.\"\"\"\n",
        "  mb = min(b.shape[0], max_images)\n",
        "  if rows is None:\n",
        "    rows = int(np.ceil(np.sqrt(mb)))\n",
        "    cols = rows\n",
        "  diff = rows * cols - mb\n",
        "  b = np.vstack([b[:mb], np.zeros([diff, b.shape[1], b.shape[2], b.shape[3]])])\n",
        "  tmp = b.reshape(-1, cols * b.shape[1], b.shape[2], b.shape[3])\n",
        "  img = np.hstack(tmp[i] for i in range(rows))\n",
        "  return img"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3dzVKRipxoz8",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# A list of attributes from which to condition generation\n",
        "# Each list element corresponds to a different fully-speciffied condition\n",
        "\n",
        "cond_attr_list = [\n",
        "    [\n",
        "        (0, 'Bald'),\n",
        "        (0, 'Black_Hair'),\n",
        "        (1, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (0, 'Male'),\n",
        "        (1, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (0, 'Bald'),\n",
        "        (0, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (1, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (0, 'Male'),\n",
        "        (1, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (0, 'Bald'),\n",
        "        (1, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (0, 'Male'),\n",
        "        (1, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (0, 'Bald'),\n",
        "        (1, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (1, 'Male'),\n",
        "        (1, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (0, 'Bald'),\n",
        "        (1, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (1, 'Male'),\n",
        "        (0, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (0, 'Bald'),\n",
        "        (1, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (1, 'Eyeglasses'),\n",
        "        (1, 'Male'),\n",
        "        (0, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (1, 'Bald'),\n",
        "        (0, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (1, 'Male'),\n",
        "        (0, 'No_Beard'),\n",
        "        (1, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (1, 'Young'),\n",
        "    ],\n",
        "    [\n",
        "        (1, 'Bald'),\n",
        "        (0, 'Black_Hair'),\n",
        "        (0, 'Blond_Hair'),\n",
        "        (0, 'Brown_Hair'),\n",
        "        (0, 'Eyeglasses'),\n",
        "        (1, 'Male'),\n",
        "        (0, 'No_Beard'),\n",
        "        (0, 'Smiling'),\n",
        "        (0, 'Wearing_Hat'),\n",
        "        (0, 'Young'),\n",
        "    ],\n",
        "]\n",
        "    \n",
        "  \n",
        "cond_attrs = []\n",
        "for attrs in cond_attr_list:\n",
        "  cond_attrs.append( np.repeat(np.array([[a[0] for a in attrs]]).astype(np.int32), m.batch_size, axis=0) )"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "9ZDz8SwKvwLf",
        "colab_type": "text"
      },
      "source": [
        "# VAE Reconstructions"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "WZ75KBiNiflS",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load D\n",
        "ckpt = os.path.join(basepath, 'D_d_2_conditional_399000.ckpt')\n",
        "m.d_saver.restore(sess, ckpt)\n",
        "# Load G\n",
        "ckpt = os.path.join(basepath, 'G_d_2_conditional_399000.ckpt')\n",
        "m.g_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sCfDDae4lV27",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "z_prior = np.random.randn(9, m.n_latent)\n",
        "labels = attr_eval[:9]\n",
        "\n",
        "ckpt = os.path.join(basepath, 'vae_best_celeba_0_crop128_beta1_xsigma1.ckpt')\n",
        "m.vae_saver.restore(sess, ckpt)\n",
        "\n",
        "z_eval_blurry = (eval_mu_xsigma1[:9] + \n",
        "                 eval_sigma_xsigma1[:9] * np.random.randn(9, m.n_latent))\n",
        "b_eval_recon_blurry = sess.run(m.x_mean, {m.z:z_eval_blurry, m.labels:np.zeros([m.batch_size, m.n_labels])})\n",
        "b_samples_blurry = sess.run(m.x_mean, {m.z:z_prior})\n",
        "\n",
        "\n",
        "ckpt = os.path.join(basepath, 'vae_best_celeba_0_crop128_beta1.ckpt')\n",
        "m.vae_saver.restore(sess, ckpt)\n",
        "\n",
        "z_eval_sharp = (eval_mu[:9] + \n",
        "                eval_sigma[:9] * np.random.randn(9, m.n_latent))\n",
        "b_eval_recon_sharp = sess.run(m.x_mean, {m.z:z_eval_sharp, m.labels:np.zeros([m.batch_size, m.n_labels])})\n",
        "b_samples_sharp = sess.run(m.x_mean, {m.z:z_prior})\n",
        "b_samples_refined = sess.run(m.x_mean, {m.q_z_sample:z_prior, \n",
        "                                        m.amortize:True, \n",
        "                                        m.labels:labels})"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Wa-YmS3csy-S",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Visualize Reconstructions\n",
        "tot = 6\n",
        "row = 5\n",
        "plt.figure(figsize=[10, 10])\n",
        "for i in range(row):\n",
        "  plt.subplot(tot, row, 1 +  i+5*1)\n",
        "  im(b_eval_recon_blurry[i])\n",
        "  plt.title('Recon')\n",
        "  plt.subplot(tot, row, 1 +  i+5*2)\n",
        "  im(b_samples_blurry[i])\n",
        "  plt.title('Sample')\n",
        "  plt.subplot(tot, row, 1 +  i+5*3)\n",
        "  im(b_eval_recon_sharp[i])\n",
        "  plt.title('Recon')\n",
        "  plt.subplot(tot, row, 1 + i+5*4)\n",
        "  im(b_samples_sharp[i])\n",
        "  plt.title('Sample')\n",
        "  plt.subplot(tot, row, 1 + i+5*5)\n",
        "  im(b_samples_refined[i])\n",
        "  plt.title('Refinement')"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "uknXc1Izxhtd",
        "colab_type": "text"
      },
      "source": [
        "# Conditional Generation "
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "BZRD_7m1bjN8",
        "colab_type": "text"
      },
      "source": [
        "## Z-Penalty = 0.0"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "KPBm_62ibfBn",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          },
          "output_extras": [
            {}
          ],
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "86ae898f-9fdc-498a-fe35-172d37015d59",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1510338488489,
          "user_tz": 480,
          "elapsed": 560,
          "user": {
            "displayName": "Jesse Engel",
            "photoUrl": "//lh3.googleusercontent.com/-MwnSfMqhHDc/AAAAAAAAAAI/AAAAAAAAACw/rJ6qJx232BA/s50-c-k-no/photo.jpg",
            "userId": "109537677218401981034"
          }
        }
      },
      "source": [
        "# Load D\n",
        "ckpt = os.path.join(basepath, 'D_d_2_conditional_399000.ckpt')\n",
        "m.d_saver.restore(sess, ckpt)\n",
        "# Load G\n",
        "ckpt = os.path.join(basepath, 'G_d_2_conditional_399000.ckpt')\n",
        "m.g_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Restoring parameters from /content/latent_constraints/D_d_2_conditional_399000.ckpt\n",
            "INFO:tensorflow:Restoring parameters from /content/latent_constraints/G_d_2_conditional_399000.ckpt\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "U0uoykd9x_d6",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Compute the Conditional Samples\n",
        "z_original = sess.run(m.prior_sample)\n",
        "z_new = [z_original]\n",
        "b_new = [sess.run(m.x_mean, {m.z:z_original})]\n",
        "\n",
        "for cond_attr in cond_attrs:\n",
        "  z_new.append(sess.run(m.z, {m.q_z_sample: z_original, m.amortize:True, m.labels:cond_attr}))\n",
        "  b_new.append(sess.run(m.x_mean, {m.z:z_new[-1]}))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "gXrfNYwS9ULk",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Plot them\n",
        "idxs = range(10)\n",
        "plt.figure(figsize=(12, 14))\n",
        "n_b = len(b_new)\n",
        "tot = 6\n",
        "barr = np.array(b_new)\n",
        "barr = np.swapaxes(barr, 0, 1)\n",
        "barr = barr[idxs, :, :, :, :]\n",
        "\n",
        "plt.figure(figsize=(18, 12))\n",
        "n_b = len(b_new)\n",
        "tot = 16\n",
        "for i, b in enumerate(b_new):\n",
        "  plt.subplot(n_b, 1, i + 1)\n",
        "  im(batch_image(b, max_images=tot, rows=tot, cols=1))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "2JbcuxIwbulE",
        "colab_type": "text"
      },
      "source": [
        "## Z-Penalty = 0.1"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "ji3F5ZzRbulG",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load D\n",
        "ckpt = os.path.join(basepath, 'D_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.d_saver.restore(sess, ckpt)\n",
        "# Load G\n",
        "ckpt = os.path.join(basepath, 'G_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.g_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "A6hyGecZbulI",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Compute the Conditional Samples\n",
        "z_original = sess.run(m.prior_sample)\n",
        "z_new = [z_original]\n",
        "b_new = [sess.run(m.x_mean, {m.z:z_original})]\n",
        "\n",
        "for cond_attr in cond_attrs:\n",
        "  z_new.append(sess.run(m.z, {m.q_z_sample: z_original, m.amortize:True, m.labels:cond_attr}))\n",
        "  b_new.append(sess.run(m.x_mean, {m.z:z_new[-1]}))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "RB_GU8OfbulL",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Plot them\n",
        "idxs = range(10)\n",
        "plt.figure(figsize=(12, 14))\n",
        "n_b = len(b_new)\n",
        "tot = 6\n",
        "barr = np.array(b_new)\n",
        "barr = np.swapaxes(barr, 0, 1)\n",
        "barr = barr[idxs, :, :, :, :]\n",
        "\n",
        "plt.figure(figsize=(18, 12))\n",
        "n_b = len(b_new)\n",
        "tot = 16\n",
        "for i, b in enumerate(b_new):\n",
        "  plt.subplot(n_b, 1, i + 1)\n",
        "  im(batch_image(b, max_images=tot, rows=tot, cols=1))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "w4vLHT2Ti_ch",
        "colab_type": "text"
      },
      "source": [
        "# Identity Preserving Transformations"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "Tx0z8MWDb-U-",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load D\n",
        "ckpt = os.path.join(basepath, 'D_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.d_saver.restore(sess, ckpt)\n",
        "# Load G\n",
        "ckpt = os.path.join(basepath, 'G_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.g_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "MrNRr_xAX943",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "def transform(z_original, \n",
        "              labels,\n",
        "              z0=None,\n",
        "              lr=1e-1, \n",
        "              n_opt=100, \n",
        "              penalty_weight=0.0,\n",
        "              r_weight=1.0, \n",
        "              attr_weight=1.0, \n",
        "              attr_weights=np.ones([1, m.n_labels]),\n",
        "              r_threshold = 0.9,\n",
        "              attr_threshold = 0.9,\n",
        "              adaptive=False,\n",
        "             ):\n",
        "\n",
        "  if z0 is None:\n",
        "    z0 = z_original\n",
        "  _ = sess.run(tf.assign(m.z_prime, z_original))\n",
        "  z_new = np.zeros([m.half_batch, m.n_latent])\n",
        "  i_threshold = np.zeros([m.half_batch])\n",
        "  z_trace = []\n",
        "  for i in range(n_opt):\n",
        "    res = sess.run([m.train_transform, \n",
        "                    m.loss_transform, \n",
        "                    m.transform_penalty,\n",
        "                    m.z_prime,\n",
        "                    m.r_pred,\n",
        "                    m.pred_attr], \n",
        "                   {m.z0: z0,\n",
        "                    m.r: np.ones([m.half_batch, 1]),\n",
        "                    m.labels: labels,\n",
        "                    m.transform: True, \n",
        "                    m.transform_lr: lr,\n",
        "                    m.transform_penalty_weight: penalty_weight,\n",
        "                    m.transform_r_weight: r_weight,\n",
        "                    m.transform_attr_weight: attr_weight,\n",
        "                    m.attr_weights: attr_weights,\n",
        "                    m.x: np.zeros([m.half_batch, m.img_width, m.img_width, 3])\n",
        "                   })\n",
        "\n",
        "    z_prime = res[3]\n",
        "    z_trace.append(z_prime)\n",
        "    pred_r = res[-2]\n",
        "    pred_attr = res[-1]\n",
        "    attr_acc = 1.0 - np.mean(attr_weights * np.abs(labels - pred_attr), axis=1)\n",
        "    transform_penalty = np.mean(res[2])\n",
        "    check_idx = np.where(i_threshold == 0)[0]\n",
        "    if len(check_idx) == 0:\n",
        "      break\n",
        "    for idx in check_idx:\n",
        "      if pred_r[idx] > r_threshold and attr_acc[idx] > attr_threshold:\n",
        "        z_new[idx] = z_prime[idx]\n",
        "        i_threshold[idx] = i\n",
        "    if adaptive:\n",
        "      r_weight = 1 - np.mean(pred_r)\n",
        "      attr_weight = 1 - np.mean(attr_acc)\n",
        "\n",
        "    if i % 100 == 1:\n",
        "      print( 'Step %d, NotConverged: %d, Loss: %0.3e, Penalty: %0.3f, '\n",
        "            'r: %0.3f, r_min: %0.3f, '\n",
        "            'attr: %0.3f, attr_min:%0.3f, ' % (\n",
        "                i,\n",
        "                len(check_idx), \n",
        "                res[1], \n",
        "                transform_penalty, \n",
        "                np.mean(pred_r), \n",
        "                np.min(pred_r),\n",
        "                np.mean(attr_acc), \n",
        "                np.min(attr_acc),\n",
        "           ))\n",
        "      \n",
        "\n",
        "  check_idx = np.where(i_threshold == 0)[0]\n",
        "  print('%d did not converge' % len(check_idx))\n",
        "  for idx in check_idx:\n",
        "    z_new[idx] = z_prime[idx]\n",
        "  return z_new, i_threshold, z_trace\n",
        "  "
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Xq4sz-6bi9OU",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# BETA1 TRANSFORMATION\n",
        "z_original = (eval_mu[:m.half_batch] + \n",
        "              eval_sigma[:m.half_batch] * np.random.randn(m.half_batch, m.n_latent))\n",
        "z = z_original\n",
        "\n",
        "label_list = (\n",
        "    (1, 1, 1e-2, 1e-3, 0.05), \n",
        "    (2, 1, 1e-2, 1e-3, 0.05), \n",
        "    (3, 1, 1e-2, 1e-3, 0.05), \n",
        "    (4, 1, 1e-2, 1e-5, 0.05), \n",
        "    (5, 0, 1e-2, 1e-3, 0.05), \n",
        "    (5, 1, 1e-2, 1e-4, 0.05), \n",
        "    (6, 0, 5e-3, 1e-4, 0.005), \n",
        "    (7, 0, 1e-2, 1e-3, 0.05), \n",
        "    (7, 1, 1e-2, 1e-3, 0.05), \n",
        "    (9, 0, 3e-2, 1e-4, 0.01), \n",
        "    (9, 1, 3e-2, 1e-4, 0.01)\n",
        ")\n",
        "z_list = [z_original]\n",
        "b_list = [sess.run(m.x_mean, {m.z: z_original})]\n",
        "\n",
        "\n",
        "for attr, value, lr, r_thresh, r_weight in label_list:\n",
        "  print('Label: %d, Value: %d' % (attr, value))\n",
        "  attr_weights=np.zeros([1, m.n_labels])\n",
        "  labels = attr_eval[:m.half_batch].copy()\n",
        "  labels[:, attr] = value\n",
        "  attr_weights[:, attr] = 1\n",
        "\n",
        "  z, i_threshold, z_trace = transform(\n",
        "      z_original, \n",
        "      labels, \n",
        "      lr=lr,\n",
        "      n_opt=300,\n",
        "      r_weight=r_weight,\n",
        "      attr_weight=1.0,\n",
        "      attr_weights=attr_weights,\n",
        "      r_threshold=0.3,\n",
        "      attr_threshold=(1 - r_thresh),\n",
        "  )\n",
        "  z_list.append(z)\n",
        "  b_list.append(sess.run(m.x_mean, {m.z: z}))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "C5LzeBfmi8g8",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "labels = attr_eval[:m.half_batch].copy()\n",
        "idxs = (0, 2, 5, 10, 7, 13,  3,)\n",
        "labels = labels[idxs, :]\n",
        "\n",
        "plt.figure(figsize=(24, 18))\n",
        "n_b = len(b_list)\n",
        "tot = 6\n",
        "barr = np.array(b_list)\n",
        "barr = np.swapaxes(barr, 0, 1)\n",
        "barr = barr[idxs, :, :, :, :]\n",
        "for i, b in enumerate(barr):\n",
        "  for j, (k, v, _, _, _) in enumerate(label_list):\n",
        "    if labels[i, k] == v:\n",
        "      barr[i, j+1] = 0.\n",
        "\n",
        "for i, b in enumerate(barr):\n",
        "  plt.subplot(12, 1, i + 1)\n",
        "  im(batch_image(b, max_images=n_b, rows=n_b, cols=1))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3B8NAeu6hEhB",
        "colab_type": "text"
      },
      "source": [
        "# Attribute Classification Accuracy\n"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "rAYzC4dTp397",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Load D\n",
        "ckpt = os.path.join(basepath, 'D_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.d_saver.restore(sess, ckpt)\n",
        "# Load G\n",
        "ckpt = os.path.join(basepath, 'G_d_2_conditional_penalty1e-1_399000.ckpt')\n",
        "m.g_saver.restore(sess, ckpt)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "k4LnDodY5YYc",
        "colab_type": "text"
      },
      "source": [
        "## Original Images"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "LEWNydfChEhG",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "tf.train.AdamOptimizer()\n",
        "batch_size = 256\n",
        "# Prediction Accuracy\n",
        "train_pred = []\n",
        "eval_pred = []\n",
        "test_pred = []\n",
        "\n",
        "for i in range(n_train / 10 /batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  batch_images = train_data[start:end]\n",
        "\n",
        "  start = i * batch_size\n",
        "  end = start + batch_size\n",
        "  res = sess.run([m.pred_classifier], {m.x: batch_images})\n",
        "  train_pred.append(res[0])\n",
        "train_pred = np.vstack(train_pred)\n",
        "\n",
        "for i in range(n_eval/batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  batch_images = eval_data[start:end]\n",
        "\n",
        "  res = sess.run([m.pred_classifier], {m.x: batch_images})\n",
        "  eval_pred.append(res[0])\n",
        "eval_pred = np.vstack(eval_pred)\n",
        "\n",
        "train_acc = (train_pred > 0.5) == attr_train[:train_pred.shape[0]]\n",
        "eval_acc = (eval_pred > 0.5) == attr_eval[:eval_pred.shape[0]]\n",
        "\n",
        "print \"Train Accuracy: %.4f\" % (np.mean(train_acc) * 100)\n",
        "print \"Eval Accuracy: %.4f\" % (np.mean(eval_acc) * 100)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "DPKFZyay5BSa",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "y_true = attr_eval[:eval_pred.shape[0]]\n",
        "y_pred = eval_pred >= 0.5\n",
        "\n",
        "prec, recall, f1, support = sklearn.metrics.classification.precision_recall_fscore_support(y_true, y_pred, )\n",
        "report = sklearn.metrics.classification_report(y_true, y_pred, digits=3, target_names=attribute_names)\n",
        "print(report)\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "EoDPpjRu5dhP",
        "colab_type": "text"
      },
      "source": [
        "## Conditional Generation"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "BJCl0ZR75Lz4",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Prediction Accuracy\n",
        "train_pred = []\n",
        "eval_pred = []\n",
        "test_pred = []\n",
        "\n",
        "for i in range(n_eval/batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  labels = attr_eval[start:end]\n",
        "  batch_z = np.random.randn(labels.shape[0], config[\"n_latent\"])\n",
        "  xtmp = np.zeros([labels.shape[0], img_width*img_width])\n",
        "  xsamp, z_prime = sess.run([m.x_mean, m.z],\n",
        "                  {m.q_z_sample: batch_z, \n",
        "                   m.amortize:True, \n",
        "                   m.labels:labels})\n",
        "  res = sess.run([m.pred_classifier], {m.x: xsamp})\n",
        "  eval_pred.append(res[0])\n",
        "eval_pred = np.vstack(eval_pred)\n",
        "\n",
        "eval_acc = (eval_pred > 0.5) == attr_eval[:eval_pred.shape[0]]\n",
        "\n",
        "eval_softmax = eval_pred.argmax(axis=-1) == attr_eval[:eval_pred.shape[0]].argmax(axis=-1)\n",
        "\n",
        "print \"Eval Accuracy: %.4f\" % (np.mean(eval_acc) * 100)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "DHWc0SO65Lz_",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "y_true = attr_eval[:eval_pred.shape[0]]\n",
        "y_pred = eval_pred >= 0.5\n",
        "\n",
        "prec, recall, f1, support = sklearn.metrics.classification.precision_recall_fscore_support(y_true, y_pred, )\n",
        "report = sklearn.metrics.classification_report(y_true, y_pred, digits=3, target_names=attribute_names)\n",
        "print(report)\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "bzvdAY0Fqg2r",
        "colab_type": "text"
      },
      "source": [
        "# TRAINING MODELS\n",
        "\n",
        "This code is for demonstration purposes. \n",
        "\n",
        "Training of the VAE and D and G can take about a day on a TitanX GPU.\n",
        "\n",
        "To train models from scratch, you will need to download the entire [CelebA dataset](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) to a base directory (assumed as `~/Desktop/CelebA/`). The steps for preprocessing the data are provided below.\n",
        "\n",
        "(img_align_celeba/, list_attr_celeba.txt, list_eval_partition.txt)"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "wrzODWr-qn-S",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Running Average\n",
        "running_N = 100\n",
        "running_N_eval = 10\n",
        "rmean = lambda data: np.mean(data[-running_N:])\n",
        "rmeane = lambda data: np.mean(data[-running_N_eval:])\n",
        "batch_size = 256"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "s8hXcrCreWU6",
        "colab_type": "text"
      },
      "source": [
        "## Prepare the Data (Crop and Pad)"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "lWcGGw9V6U3L",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "basepath = os.path.expanduser('~/Desktop/CelebA/')\n",
        "save_path = basepath"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1m0a42lEmf3W",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "partition = np.loadtxt(basepath + 'list_eval_partition.txt', usecols=(1,))\n",
        "train_mask = (partition == 0)\n",
        "eval_mask = (partition == 1)\n",
        "test_mask = (partition == 2)\n",
        "\n",
        "print \"Train: %d, Validation: %d, Test: %d, Total: %d\" % (train_mask.sum(), eval_mask.sum(), test_mask.sum(), partition.shape[0])"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_rVovMjEm75q",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "attributes = pd.read_table(basepath + 'list_attr_celeba.txt', skiprows=1, delim_whitespace=True, usecols=range(1, 41))\n",
        "attribute_names = attributes.columns.values\n",
        "attribute_values = attributes.values"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "le9fdx3itF_o",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "attr_train = attribute_values[train_mask]\n",
        "attr_eval = attribute_values[eval_mask]\n",
        "attr_test = attribute_values[test_mask]\n",
        "\n",
        "attr_train[attr_train == -1] = 0\n",
        "attr_eval[attr_eval == -1] = 0\n",
        "attr_test[attr_test == -1] = 0\n",
        "\n",
        "np.save(basepath + 'attr_train.npy', attr_train)\n",
        "np.save(basepath + 'attr_eval.npy', attr_eval)\n",
        "np.save(basepath + 'attr_test.npy', attr_test)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Kb71ejrxO3Rn",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "def pil_crop_downsample(x, width, out_width):\n",
        "  half_shape = tuple((i - width) / 2 for i in x.size)\n",
        "  x = x.crop([half_shape[0], half_shape[1], half_shape[0] + width, half_shape[1] + width])\n",
        "  return x.resize([out_width, out_width], resample=PIL.Image.ANTIALIAS)\n",
        "\n",
        "def load_and_adjust_file(filename, width, outwidth):\n",
        "  img = PIL.Image.open(filename)\n",
        "  img = pil_crop_downsample(img, width, outwidth)\n",
        "  img = np.array(img, np.float32) / 255.\n",
        "  return img"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "5Ikg2GHwgrv6",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# CELEBA images are (218 x 178) originally\n",
        "filenames = np.sort(glob(basepath + 'img_align_celeba/*.jpg'))\n",
        "\n",
        "crop_width = 128\n",
        "img_width = 64\n",
        "postfix = '_crop_%d_res_%d.npy' % (crop_width, img_width)\n",
        "\n",
        "n_files = len(filenames)\n",
        "all_data = np.zeros([n_files, img_width, img_width, 3], np.float32)\n",
        "for i, fname in enumerate(filenames):\n",
        "  all_data[i, :, :] = load_and_adjust_file(fname, crop_width, img_width)\n",
        "  if i % 10000 == 0:\n",
        "    print('%.2f percent done' % (float(i)/n_files * 100.0))\n",
        "train_data = all_data[train_mask]\n",
        "eval_data = all_data[eval_mask]\n",
        "test_data = all_data[test_mask]\n",
        "np.save(basepath + 'train' + postfix, train_data)\n",
        "np.save(basepath + 'eval' + postfix, eval_data)\n",
        "np.save(basepath + 'test' + postfix, test_data)\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "i11MFfZ4sA2Y",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Train the VAE"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "AWCOid2F_B2v",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "sess.run(tf.variables_initializer(var_list=m.vae_vars))\n",
        "\n",
        "# Train the VAE\n",
        "results = []\n",
        "results_eval = []\n",
        "\n",
        "traces = {'i': [],\n",
        "          'i_eval': [],\n",
        "          'loss': [],\n",
        "          'loss_eval': [],\n",
        "          'recons': [],\n",
        "          'recons_eval': [],\n",
        "          'kl': [],\n",
        "          'kl_eval': []}\n",
        "\n",
        "n_iters = 200000\n",
        "vae_lr_ = np.logspace(np.log10(3e-4), np.log10(1e-6), n_iters)\n",
        "\n",
        "for i in range(n_iters):\n",
        "  start = (i * batch_size) % n_train\n",
        "  end = start + batch_size\n",
        "  batch = train_data[start:end]\n",
        "\n",
        "  res = sess.run([m.train_vae, \n",
        "                  m.vae_loss, \n",
        "                  m.mean_recons, \n",
        "                  m.mean_KL], \n",
        "                 {m.x: batch,\n",
        "                  m.vae_lr: vae_lr_[i],\n",
        "                  m.amortize: False,\n",
        "                  m.labels: attr_train[start:end]})\n",
        "  \n",
        "  traces['loss'].append(res[1])\n",
        "  traces['recons'].append(res[2])\n",
        "  traces['kl'].append(res[3])\n",
        "  traces['i'].append(i)\n",
        "\n",
        "  if i % 10 == 0:\n",
        "    start = (i * batch_size) % n_eval\n",
        "    end = start + batch_size\n",
        "    batch = eval_data[start:end]\n",
        "    res_eval = sess.run([m.vae_loss, m.mean_recons, m.mean_KL], \n",
        "                        {m.x: batch, m.labels: attr_eval[start:end]})\n",
        "    traces['loss_eval'].append(res_eval[0])\n",
        "    traces['recons_eval'].append(res_eval[1])\n",
        "    traces['kl_eval'].append(res_eval[2])\n",
        "    traces['i_eval'].append(i)\n",
        "\n",
        "    print('Step %5d \\t TRAIN \\t Loss: %0.3f, Recon: %0.3f, KL: %0.3f '\n",
        "          '\\t EVAL \\t  Loss: %0.3f, Recon: %0.3f, KL: %0.3f' % (i, \n",
        "                                                                rmean(traces['loss']), \n",
        "                                                                rmean(traces['recons']), \n",
        "                                                                rmean(traces['kl']), \n",
        "                                                                rmeane(traces['loss_eval']), \n",
        "                                                                rmeane(traces['recons_eval']), \n",
        "                                                                rmeane(traces['kl_eval']) ))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "u_dIVjWl_QCb",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "plt.figure(figsize=(18,6))\n",
        "\n",
        "plt.subplot(131)\n",
        "plt.plot(traces['i'], traces['loss'])\n",
        "plt.plot(traces['i_eval'], traces['loss_eval'])\n",
        "plt.title('Loss')\n",
        "# plt.ylim(30, 100)\n",
        "\n",
        "plt.subplot(132)\n",
        "plt.plot(traces['i'], traces['recons'])\n",
        "plt.plot(traces['i_eval'], traces['recons_eval'])\n",
        "plt.title('Recons')\n",
        "# plt.ylim(-100, -30)\n",
        "\n",
        "plt.subplot(133)\n",
        "plt.plot(traces['i'], traces['kl'])\n",
        "plt.plot(traces['i_eval'], traces['kl_eval'])\n",
        "plt.title('KL')\n",
        "# plt.ylim(10, 100)\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "T5zLm-rSuFQm",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Train D and G jointly"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "J1hT5rNSTtOW",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Precompute means and vars\n",
        "train_mu = []\n",
        "train_sigma = []\n",
        "n_batches = int(np.ceil(float(n_train) / batch_size))\n",
        "for i in range(n_batches):\n",
        "  if i % 1000 == 0:\n",
        "    print '%.1f Done' % (float(i) / n_train * batch_size * 100)\n",
        "  start = i * batch_size\n",
        "  end = start + batch_size\n",
        "  res = sess.run([m.mu, m.sigma], {m.x: train_data[start:end]})\n",
        "  train_mu.append(res[0])\n",
        "  train_sigma.append(res[1])\n",
        "train_mu = np.vstack(train_mu)\n",
        "train_sigma = np.vstack(train_sigma)\n",
        "sigma_mean = train_sigma.mean(0, keepdims=True)\n",
        "print train_mu.shape, train_sigma.shape, train_data.shape"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "4-oILpXELVv8",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Precompute means and vars\n",
        "eval_mu = []\n",
        "eval_sigma = []\n",
        "n_batches = int(np.ceil(float(n_eval) / batch_size))\n",
        "for i in range(n_batches):\n",
        "  if i % 1000 == 0:\n",
        "    print '%.1f Done' % (float(i) / n_eval * batch_size * 100)\n",
        "  start = i * batch_size\n",
        "  end = start + batch_size\n",
        "  res = sess.run([m.mu, m.sigma], {m.x: eval_data[start:end]})\n",
        "  eval_mu.append(res[0])\n",
        "  eval_sigma.append(res[1])\n",
        "eval_mu = np.vstack(eval_mu)\n",
        "eval_sigma = np.vstack(eval_sigma)\n",
        "sigma_mean_eval = eval_sigma.mean(0, keepdims=True)\n",
        "print eval_mu.shape, eval_sigma.shape, eval_data.shape"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "tycqyI5-Z1DH",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "slide"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "plt.plot(np.sort(sigma_mean.flatten()))\n",
        "plt.plot(np.sort(sigma_mean_eval.flatten()))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "qbYzYGh-yZgR",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "slide"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# With eval loss, looking for overfitting\n",
        "sess.run(tf.variables_initializer(var_list=m.d_vars))\n",
        "sess.run(tf.variables_initializer(var_list=m.g_vars))\n",
        "\n",
        "# Declare hyperparameters\n",
        "results_D = []\n",
        "results_G = []\n",
        "traces = {'i': [], \n",
        "          'i_pred': [], \n",
        "          'D': [], \n",
        "          'G': [], \n",
        "          'G_real': [],\n",
        "          'g_penalty': [],\n",
        "          'pred_train': [], \n",
        "          'pred_eval': [], \n",
        "          'pred_prior': [], \n",
        "          'pred_gen': [],\n",
        "          'z_dist_eval': [],\n",
        "          'attr_loss': [],\n",
        "          'attr_acc': [],\n",
        "         }\n",
        "\n",
        "\n",
        "n_iters = 200000\n",
        "d_lr_ = np.logspace(-4, -4, n_iters)\n",
        "g_lr_ = np.logspace(-4, -4, n_iters)\n",
        "\n",
        "g_penalty_weight_ = 0.1\n",
        "lambda_weight_ = 10\n",
        "z_sigma_mean_ = sigma_mean\n",
        "\n",
        "percentage_prior_fake = 0.1\n",
        "N_between_update_G = 10\n",
        "N_between_eval = 100\n",
        "\n",
        "n_train = train_mu.shape[0]\n",
        "n_eval = eval_mu.shape[0]\n",
        "\n",
        "# Training Loop\n",
        "for i in range(n_iters):\n",
        "  start = (i * batch_size/2) % n_train\n",
        "  end = start + batch_size/2\n",
        "  fake_start = np.random.choice(np.arange(n_train - batch_size/2))\n",
        "  fake_end = fake_start + batch_size/2\n",
        "  \n",
        "  real_img = train_data[start:end]\n",
        "  n_batch = real_img.shape[0]\n",
        "  if n_batch == batch_size/2 and start != fake_start:\n",
        "    # Compare real vs. fake\n",
        "    fake_z_prior = np.random.randn(batch_size/2, n_latent)\n",
        "    real_attr = attr_train[start:end].astype(np.int32)\n",
        "    fake_attr = attr_train[fake_start:fake_end].astype(np.int32)            \n",
        "    real_z = train_mu[start:end] + train_sigma[start:end] * np.random.randn(batch_size/2, n_latent)\n",
        "\n",
        "\n",
        "    \n",
        "    if np.random.rand(1) < percentage_prior_fake:\n",
        "      # Use Prior for fake_samples      \n",
        "      all_z = np.vstack([real_z, fake_z_prior, real_z])\n",
        "    else:\n",
        "      # Use Generator to make fake_samples \n",
        "      fake_z_gen = sess.run(m.z, {m.q_z_sample: fake_z_prior, \n",
        "                                  m.amortize:True, \n",
        "                                  m.labels: real_attr,})      \n",
        "      all_z = np.vstack([real_z, fake_z_gen, real_z])\n",
        "    all_attr = np.vstack([real_attr, real_attr, fake_attr]) \n",
        "\n",
        "      \n",
        "    # Train Discriminator\n",
        "    real_r = np.ones([batch_size/2, 1])\n",
        "    fake_r = np.zeros([batch_size/2, 1])\n",
        "    all_r = np.concatenate([real_r, fake_r, fake_r])\n",
        "    res_d = sess.run([m.train_d, m.d_loss], {m.z: all_z, \n",
        "                                         m.r: all_r,\n",
        "                                         m.d_lr: d_lr_[i],\n",
        "                                         m.lambda_weight: lambda_weight_,\n",
        "                                         m.labels: all_attr,})\n",
        "    \n",
        "    # Train Generator\n",
        "    if i % N_between_update_G == 0:\n",
        "      if g_penalty_weight_ > 0:\n",
        "        # Train on real data\n",
        "        res_g_real = sess.run([m.train_g, m.g_loss, m.g_penalty], \n",
        "                              {m.q_z_sample: real_z,\n",
        "                               m.amortize: True,\n",
        "                               m.g_penalty_weight: g_penalty_weight_,\n",
        "                               m.z_sigma_mean: z_sigma_mean_,\n",
        "                               m.g_lr: g_lr_[i],\n",
        "                               m.labels: real_attr,})\n",
        "        traces['G_real'].append(res_g_real[1])\n",
        "\n",
        "      # Train on generated data\n",
        "      res_g = sess.run([m.train_g, m.g_loss, m.g_penalty], \n",
        "                       {m.q_z_sample: fake_z_prior,\n",
        "                        m.amortize: True,\n",
        "                        m.g_penalty_weight: g_penalty_weight_,\n",
        "                        m.z_sigma_mean: z_sigma_mean_,\n",
        "                        m.g_lr: g_lr_[i],\n",
        "                        m.labels: real_attr,})\n",
        "\n",
        "\n",
        "\n",
        "    traces['i'].append(i)\n",
        "    traces['D'].append(res_d[1])\n",
        "    traces['G'].append(res_g[1])\n",
        "    traces['g_penalty'].append(res_g[2])\n",
        "\n",
        "    if i % N_between_eval == 0:\n",
        "      eval_start = np.random.choice(np.arange(n_eval - batch_size/2))\n",
        "      eval_end = eval_start + batch_size/2\n",
        "      real_attr_eval = attr_eval[eval_start:eval_end].astype(np.int32)\n",
        "      real_z_eval = eval_mu[eval_start:eval_end] + eval_sigma[eval_start:eval_end] * np.random.randn(batch_size/2, n_latent)\n",
        "      z_eval_gen = sess.run(m.z, {m.q_z_sample: real_z_eval, \n",
        "                                  m.amortize:True, \n",
        "                                  m.labels: real_attr,})      \n",
        "      fake_z_gen = sess.run(m.z, {m.q_z_sample: fake_z_prior, \n",
        "                                  m.amortize:True, \n",
        "                                  m.labels: real_attr,})      \n",
        "      \n",
        "      pred_train_ = np.mean(sess.run([m.r_pred], {m.z: real_z, m.labels: real_attr}))\n",
        "      pred_eval_ = np.mean(sess.run([m.r_pred], {m.z: real_z_eval, m.labels: real_attr_eval}))\n",
        "      pred_prior_ = np.mean(sess.run([m.r_pred], {m.z: fake_z_prior, m.labels: real_attr}))\n",
        "      pred_gen_ = np.mean(sess.run([m.r_pred], {m.z: fake_z_gen, m.labels: real_attr}))\n",
        "\n",
        "      traces['i_pred'].append(i)\n",
        "      traces['pred_train'].append(pred_train_)\n",
        "      traces['pred_eval'].append(pred_eval_)\n",
        "      traces['pred_prior'].append(pred_prior_)\n",
        "      traces['pred_gen'].append(pred_gen_)\n",
        "      traces['z_dist_eval'].append(np.mean(((z_eval_gen - real_z_eval)/z_sigma_mean_)**2))\n",
        "      print 'PRED Step %d, \\t TRAIN: %.2e \\t EVAL: %.2e \\t PRIOR: %.2e \\t GEN: %.2e ' % (i, \n",
        "                                                                                         rmeanp(traces['pred_train']),  \n",
        "                                                                                         rmeanp(traces['pred_eval']),  \n",
        "                                                                                         rmeanp(traces['pred_prior']),  \n",
        "                                                                                         rmeanp(traces['pred_gen']))\n",
        "    "
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1t80OZTfuFQx",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "slide"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "plt.figure(figsize=(18,12))\n",
        "\n",
        "plt.subplot(4, 1, 1)\n",
        "plt.plot(traces['i_pred'], traces['pred_train'], label='train')\n",
        "plt.plot(traces['i_pred'], traces['pred_eval'], label='eval')\n",
        "plt.plot(traces['i_pred'], traces['pred_prior'], label='prior')\n",
        "plt.plot(traces['i_pred'], traces['pred_gen'], label='gen')\n",
        "plt.ylabel('Predictions')\n",
        "plt.legend(loc='upper right')\n",
        "\n",
        "plt.subplot(4, 1, 2)\n",
        "plt.plot(traces['i'], traces['G'])\n",
        "plt.ylabel('G Loss')\n",
        "\n",
        "plt.subplot(4, 1, 3)\n",
        "# plt.semilogy(traces['i'], traces['D'])\n",
        "plt.plot(traces['i'], traces['D'])\n",
        "plt.ylabel('D Loss')\n",
        "\n",
        "plt.subplot(4, 1, 4)\n",
        "plt.semilogy(traces['i_pred'], traces['z_dist_eval'])\n",
        "plt.ylabel('Weighted Z Distance Eval')"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "PGiELds5QHYh",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Train Classifier"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "hP1BttVSQHYk",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "sess.run(tf.variables_initializer(var_list=m.classifier_vars))\n",
        "# Train the Classifier\n",
        "results = []\n",
        "results_eval = []\n",
        "\n",
        "running_N = 100\n",
        "running_loss = 1\n",
        "running_loss_eval = 1\n",
        "\n",
        "classifier_lr_ = 3e-4\n",
        "\n",
        "# Train\n",
        "for i in range(40000):\n",
        "  start = (i * batch_size) % n_train\n",
        "  end = start + batch_size\n",
        "  batch_images = train_data[start:end]\n",
        "  batch_labels = attr_train[start:end]\n",
        "\n",
        "  res = sess.run([m.train_classifier, \n",
        "                  m.classifier_loss], \n",
        "                 {m.x: batch_images, \n",
        "                  m.labels: batch_labels.astype(np.int32),\n",
        "                  m.classifier_lr: classifier_lr_})\n",
        "  running_loss += (res[1] - running_loss) / running_N \n",
        "  if i % 10 == 1:\n",
        "    start = (i * batch_size) % n_eval\n",
        "    end = start + batch_size\n",
        "    eval_images = eval_data[start:end]\n",
        "    eval_labels = attr_eval[start:end]\n",
        "    res_eval = sess.run([m.classifier_loss], \n",
        "                        {m.x: eval_images, \n",
        "                         m.labels: eval_labels.astype(np.int32)})\n",
        "    running_loss_eval += (res_eval[0] - running_loss_eval) / (running_N / 10)\n",
        "      \n",
        "    results.append([i] + res[1:])\n",
        "    results_eval.append([i] + res_eval[0:])\n",
        "\n",
        "  if i % 10 == 1:\n",
        "    print('Step %d, \\t TRAIN \\t Loss: %0.3f \\t EVAL \\t Loss: %0.3f' % (i, running_loss, running_loss_eval))"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "15tYzkGhQHYs",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "plot_train = np.array(results).T\n",
        "plot_eval = np.array(results_eval).T\n",
        "plt.figure(figsize=(18,6))\n",
        "plt.plot(plot_train[0],plot_train[1])\n",
        "plt.plot(plot_eval[0],plot_eval[1])\n",
        "plt.ylim(1e-1, 1)\n",
        "plt.title('Loss')"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "LrM3qk7rQHYw",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Prediction Accuracy\n",
        "train_pred = []\n",
        "eval_pred = []\n",
        "test_pred = []\n",
        "\n",
        "for i in range(n_train / 10 /batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  batch_images = train_data[start:end]\n",
        "\n",
        "  start = i * batch_size\n",
        "  end = start + batch_size\n",
        "  res = sess.run([m.pred_classifier], {m.x: batch_images})\n",
        "  train_pred.append(res[0])\n",
        "train_pred = np.vstack(train_pred)\n",
        "\n",
        "for i in range(n_eval/batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  batch_images = eval_data[start:end]\n",
        "\n",
        "  res = sess.run([m.pred_classifier], {m.x: batch_images})\n",
        "  eval_pred.append(res[0])\n",
        "eval_pred = np.vstack(eval_pred)\n",
        "\n",
        "train_acc = (train_pred > 0.5) == attr_train[:train_pred.shape[0]]\n",
        "eval_acc = (eval_pred > 0.5) == attr_eval[:eval_pred.shape[0]]\n",
        "\n",
        "print \"Train Accuracy: %.4f\" % (np.mean(train_acc) * 100)\n",
        "print \"Eval Accuracy: %.4f\" % (np.mean(eval_acc) * 100)"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "t5FXkG-_l4SQ",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "y_true = attr_eval[:eval_pred.shape[0]]\n",
        "y_pred = eval_pred >= 0.5\n",
        "\n",
        "prec, recall, f1, support = sklearn.metrics.classification.precision_recall_fscore_support(y_true, y_pred, )\n",
        "report = sklearn.metrics.classification_report(y_true, y_pred, digits=2, target_names=attribute_names)\n",
        "print report\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Lv_B7JmvqMP0",
        "colab_type": "text",
        "slideshow": {
          "slide_type": "slide"
        }
      },
      "source": [
        "# Train Attribute Classifier (D_attr) in z space"
      ],
      "cell_type": "markdown"
    },
    {
      "metadata": {
        "id": "ukJYuokLqH9s",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Train the Discriminator\n",
        "sess.run(tf.variables_initializer(var_list=m.d_attr_vars))\n",
        "traces = {\"i\": [],\n",
        "          \"i_eval\": [],\n",
        "          \"D_loss\": [],\n",
        "          \"D_loss_eval\": [],\n",
        "          \"accuracy\": [],\n",
        "         }\n",
        "running_N = 800\n",
        "running_N_eval = 80\n",
        "\n",
        "n_iters = 10000\n",
        "d_attr_lr_ = np.logspace(-4, -4, n_iters)\n",
        "\n",
        "for i in range(n_iters):\n",
        "  start = (i * batch_size/2) % n_train\n",
        "  end = min(start + batch_size/2, n_train)\n",
        "  batch = train_mu[start:end] + train_sigma[start:end] * np.random.randn(end-start, n_latent)\n",
        "  if batch.shape[0] == batch_size/2:\n",
        "    labels = attr_train[start:end]\n",
        "\n",
        "    # Train\n",
        "    res = sess.run([m.train_d_attr, m.d_loss_attr], \n",
        "                   {m.z: batch, \n",
        "                    m.labels: labels, \n",
        "                    m.d_attr_lr: d_attr_lr_[i],})\n",
        "\n",
        "    traces['i'].append(i)\n",
        "    traces['D_loss'].append(res[1])\n",
        "\n",
        "  if i % 10 == 1:\n",
        "    start = (i * batch_size/2) % n_eval\n",
        "    end = min(start + batch_size/2, n_eval)\n",
        "    batch = eval_mu[start:end] + eval_sigma[start:end] * np.random.randn(end-start, n_latent)\n",
        "    if batch.shape[0] == batch_size/2:\n",
        "      labels = attr_eval[start:end]\n",
        "\n",
        "      res_eval = sess.run([m.d_loss_attr, m.pred_attr], \n",
        "                          {m.z: batch,\n",
        "                           m.labels: labels,})\n",
        "      \n",
        "      y_true = labels\n",
        "      y_pred = (res_eval[1] > 0.5)\n",
        "      accuracy = np.mean(y_true == y_pred)\n",
        "      \n",
        "      traces['i_eval'].append(i)\n",
        "      traces['D_loss_eval'].append(res_eval[0])\n",
        "      traces['accuracy'].append(accuracy)\n",
        "      \n",
        "\n",
        "  if i % 100 == 0:\n",
        "    print 'Step %d, \\t TRAIN \\t Loss: %0.3f \\t EVAL \\t Loss: %0.3f \\t Accuracy: %0.3f' % (i, \n",
        "                                                                       np.mean(traces['D_loss'][-running_N_eval:]), \n",
        "                                                                       np.mean(traces['D_loss_eval'][-running_N_eval:]),\n",
        "                                                                       np.mean(traces['accuracy'][-running_N_eval:]), \n",
        "                                                                       )"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Pl9vrH62qH9w",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "plt.figure(figsize=(18,18))\n",
        "plt.subplot(3, 1, 1)\n",
        "plt.semilogy(traces['i'], traces['D_loss'], label='train')\n",
        "plt.semilogy(traces['i_eval'], traces['D_loss_eval'], label='eval')\n",
        "plt.ylabel('Loss')\n",
        "plt.legend(loc='upper right')\n",
        "\n",
        "plt.subplot(3, 1, 2)\n",
        "plt.plot(traces['i_eval'], traces['accuracy'], label=\"eval\")\n",
        "plt.ylabel('Prediction Accuracy')\n",
        "plt.legend(loc='upper right')"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GBYWe8RLsnAb",
        "colab_type": "code",
        "slideshow": {
          "slide_type": "-"
        },
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "# Prediction Accuracy\n",
        "train_pred = []\n",
        "eval_pred = []\n",
        "test_pred = []\n",
        "\n",
        "for i in range(n_train / 10 /batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  batch_images = train_data[start:end]\n",
        "\n",
        "  start = i * batch_size\n",
        "  end = start + batch_size\n",
        "  res = sess.run([m.pred_attr], {m.x: batch_images, \n",
        "                                 m.labels: attr_train[start:end]})\n",
        "  train_pred.append(res[0])\n",
        "train_pred = np.vstack(train_pred)\n",
        "\n",
        "for i in range(n_eval/batch_size):\n",
        "  start = (i * batch_size)\n",
        "  end = start + batch_size\n",
        "  batch_images = eval_data[start:end]\n",
        "\n",
        "  res = sess.run([m.pred_attr], {m.x: batch_images,\n",
        "                                 m.labels: attr_eval[start:end]})\n",
        "  eval_pred.append(res[0])\n",
        "eval_pred = np.vstack(eval_pred)\n",
        "\n",
        "train_acc = (train_pred > 0.5) == attr_train[:train_pred.shape[0]]\n",
        "eval_acc = (eval_pred > 0.5) == attr_eval[:eval_pred.shape[0]]\n",
        "\n",
        "print \"Train Accuracy: %.4f\" % (np.mean(train_acc) * 100)\n",
        "print \"Eval Accuracy: %.4f\" % (np.mean(eval_acc) * 100)\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "eDc6rOaNsnAi",
        "colab_type": "code",
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        }
      },
      "source": [
        "y_true = attr_eval[:eval_pred.shape[0]]\n",
        "y_pred = eval_pred >= 0.5\n",
        "\n",
        "prec, recall, f1, support = sklearn.metrics.classification.precision_recall_fscore_support(y_true, y_pred, )\n",
        "report = sklearn.metrics.classification_report(y_true, y_pred, digits=2, target_names=attribute_names)\n",
        "print(report)\n"
      ],
      "cell_type": "code",
      "execution_count": 0,
      "outputs": []
    }
  ]
}